import numpy as np
import cv2
import torch

def adjust_lr(optimizer, init_lr, gamma):
	for p in optimizer.param_groups:
		p['lr'] = init_lr * gamma
	return optimizer

def tocuda(data):
	# convert tensor data in dictionary to cuda when it is a tensor
	for key in data.keys():
		if type(data[key]) == torch.Tensor:
			data[key] = data[key].cuda()
	return data

def denorm(x, T):
    x = x * np.array([T[0,0], T[1,1]]) + np.array([T[0,2], T[1,2]])
    return x

def np_skew_symmetric(v):

    zero = np.zeros_like(v[:, 0])

    M = np.stack([
        zero, -v[:, 2], v[:, 1],
        v[:, 2], zero, -v[:, 0],
        -v[:, 1], v[:, 0], zero,
    ], axis=1)

    return M

def torch_skew_symmetric(v):

    zero = torch.zeros_like(v[:, 0])

    M = torch.stack([
        zero, -v[:, 2], v[:, 1],
        v[:, 2], zero, -v[:, 0],
        -v[:, 1], v[:, 0], zero,
    ], dim=1)

    return M

'''
这段代码定义了一个名为 estimate_pose_norm_kpts 的函数，用于估计两个图像之间的相对姿态。它接受两个关键点集 kpts0 和 kpts1，以及阈值 thresh 和置信度 conf 作为输入参数。
1. 检查关键点的数量，如果小于5，则返回 None，因为至少需要5个关键点才能估计姿态.
2. 使用 cv2.findEssentialMat 函数计算本质矩阵 E 和掩码 mask，该函数基于RANSAC方法，从关键点中筛选出内点，并计算本质矩阵.
3. 确保本质矩阵 E 不为空，否则程序将抛出异常.
4. 初始化一些变量，包括最佳内点数量 best_num_inliers、新的掩码 new_mask 和返回值 ret.
5. 将本质矩阵 E 按每3个元素分割，对每个分割后的矩阵 _E 进行处理.
6. 对于每个 _E，使用 cv2.recoverPose 函数恢复相机的姿态 R、平移向量 t 和掩码 mask_，该函数基于本质矩阵和关键点，计算相机的相对姿态.
7. 如果当前内点数量 n 大于最佳内点数量，则更新最佳内点数量和返回值 ret，其中返回值 ret 包括姿态矩阵 R、平移向量 t 和掩码 mask.ravel() > 0，后者表示内点的布尔掩码.
8. 最后，返回最佳的姿态估计结果 ret.
'''
def estimate_pose_norm_kpts(kpts0, kpts1, thresh=1e-3, conf=0.99999):
    # 如果关键点的数量小于5，则返回None
	if len(kpts0) < 5:
		return None

    # TODO 使用cv2.findEssentialMat函数计算本质矩阵E和掩码mask
    # 参数包括关键点kpts0和kpts1，单位矩阵np.eye(3)，阈值thresh，置信度prob，以及RANSAC方法
	E, mask = cv2.findEssentialMat(
	kpts0, kpts1, np.eye(3), threshold=thresh, prob=conf,
	method=cv2.RANSAC)

    # 断言：如果该语句出错就直接跳出
	assert E is not None

	best_num_inliers = 0
	new_mask = mask
	ret = None
	for _E in np.split(E, len(E) / 3):
		n, R, t, mask_ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask)
		if n > best_num_inliers:
			best_num_inliers = n
			ret = (R, t[:, 0], mask.ravel() > 0)

	return ret

def estimate_pose_from_E(kpts0, kpts1, mask, E):
    assert E is not None
    mask = mask.astype(np.uint8)
    E = E.astype(np.float64)
    kpts0 = kpts0.astype(np.float64)
    kpts1 = kpts1.astype(np.float64)
    I = np.eye(3).astype(np.float64)

    best_num_inliers = 0
    ret = None

    for _E in np.split(E, len(E) / 3):

        n, R, t, _ = cv2.recoverPose(
            _E, kpts0, kpts1, I, 1e9, mask=mask)

        if n > best_num_inliers:
            best_num_inliers = n
            ret = (R, t[:, 0], mask.ravel() > 0)
    return ret

def angle_error_mat(R1, R2):
    cos = (np.trace(np.dot(R1.T, R2)) - 1) / 2
    cos = np.clip(cos, -1., 1.)  # numercial errors can make it out of bounds
    return np.rad2deg(np.abs(np.arccos(cos)))

def angle_error_vec(v1, v2):
    n = np.linalg.norm(v1) * np.linalg.norm(v2)
    return np.rad2deg(np.arccos(np.clip(np.dot(v1, v2) / n, -1.0, 1.0)))

def compute_pose_error(T_0to1, R, t):
    R_gt = T_0to1[:3, :3]
    t_gt = T_0to1[:3, 3]
    error_t = angle_error_vec(t, t_gt)
    error_t = np.minimum(error_t, 180 - error_t)  # ambiguity of E estimation
    error_R = angle_error_mat(R, R_gt)
    return error_t, error_R

'''
这个函数 pose_auc 用于计算一个姿态估计任务中，不同阈值下的平均召回率（Area Under Curve, AUC）。具体来说，它计算在给定的误差阈值范围内，姿态估计的召回率曲线下的面积。以下是该函数的详细解释：

参数
errors：一个列表或数组，包含姿态估计的误差值。这些误差值通常是姿态估计与真实姿态之间的某种距离度量，例如旋转误差或平移误差.
thresholds：一个列表或数组，包含用于评估的误差阈值。这些阈值用于确定在不同误差范围内，姿态估计的召回率.
返回值
aucs：一个列表，包含在每个给定阈值下的AUC值.
应用场景
这个函数通常用于评估姿态估计模型的性能。通过计算不同误差阈值下的AUC，可以了解模型在不同误差范围内的召回率表现，从而全面评估模型的准确性和鲁棒性
'''
def pose_auc(errors, thresholds):
    # 获取误差值从小到大排序的索引.
    sort_idx = np.argsort(errors)
    # 根据排序索引对误差值进行排序，以便后续计算召回率.
    errors = np.array(errors.copy())[sort_idx]
    # 计算排序后的误差值对应的召回率。召回率是当前误差值在所有误差值中的排名比例，即第i个误差值的召回率为i/len(errors).
    recall = (np.arange(len(errors)) + 1) / len(errors)
    # 设置起始点
    # 在误差值数组前面添加一个0，表示误差为0时的起始点.
    errors = np.r_[0., errors]
    # 在召回率数组前面添加一个0，表示误差为0时的召回率为0.
    recall = np.r_[0., recall]
    # 计算每个阈值下的AUC：
    aucs = []
    for t in thresholds:
        # 找到误差值数组中第一个大于或等于阈值t的位置索引.
        last_index = np.searchsorted(errors, t)
        # 获取在该阈值下的召回率值，包括小于阈值的所有召回率值和最后一个小于阈值的召回率值.
        r = np.r_[recall[:last_index], recall[last_index-1]]
        # 获取在该阈值下的误差值，包括小于阈值的所有误差值和阈值本身.
        e = np.r_[errors[:last_index], t]
        # 使用梯形法则（np.trapz）计算召回率曲线下的面积，并除以阈值t，得到该阈值下的AUC值.
        aucs.append(np.trapz(r, x=e)/t)
    return aucs
